#pragma once
#include <utility>
#include <cassert>
#include <iostream>
namespace wyd
{
    template <class K, class V>
    struct TreeNode
    {
        typedef TreeNode<K, V> Self;
        std::pair<K, V> _kv;
        int _bf;
        Self *_left;
        Self *_right;
        Self *_parent;
        
        TreeNode(const std::pair<K, V>& kv)
            : _kv(kv)
            , _bf(0)
            , _left(nullptr)
            , _right(nullptr)
            , _parent(nullptr)
        {}
    };

    template<class K, class V>
    class AVLtree
    {
        typedef TreeNode<K, V> Node;
    public:
        bool insert(const std::pair<K, V>& kv);
        void isBinarySerchTree()
        {
            _isBinarySerchTree(_root);
            std::cout << std::endl;
        }
        bool isBalanceTree()
        {
            return _isBalanceTree(_root);
        }
        bool BalanceFactor()
        {
            return _BalanceFactor(_root);
        }

    private:
        void RotateL(Node* parent);
        void RotateR(Node* parent);
        void RotateRL(Node* parent);
        void RotateLR(Node* parent);
        void _isBinarySerchTree(Node* root);
        bool _isBalanceTree(Node* root);
        bool _BalanceFactor(Node* root);

        int Hight(Node* root);

    private:
        Node *_root = nullptr;
    };

    template<class K, class V>
    bool AVLtree<K, V>::insert(const std::pair<K, V>& kv)
    {
        if (_root == nullptr)
        {
            _root = new Node(kv);
            return true;
        }

        //找位置插入
        Node* parent = nullptr;
        Node* cur = _root;
        while (cur != nullptr)
        {
            parent = cur;
            if (cur->_kv.first < kv.first)
            {
                cur = cur->_left;
            }
            else if (cur->_kv.first > kv.first)
            {
                cur = cur->_right;
            }
            else // ==
                return false;
        }

        //插入
        cur = new Node(kv);
        cur->_parent = parent;
        if (parent->_kv.first < kv.first)
        {
            parent->_left = cur;
        }
        else if (parent->_kv.first > kv.first)
        {
            parent->_right = cur;
        }
        else
        {
            assert(false);
            exit(-1);
        }

        //调节平衡因子
        while (parent != nullptr)
        {
            if (parent->_left == cur)
            {
                --parent->_bf;
            }
            else
            {
                ++parent->_bf;
            }

            if (parent->_bf == 0)
            {
                break;
            }
            else if (parent->_bf == 1 || parent->_bf == -1)
            {
                //更新cur
                cur = parent;
                parent = parent->_parent;
            }
            else if (parent->_bf == 2 || parent->_bf == -2)
            {
                //旋转
                if (parent->_bf == 2)
                {
                    if (cur->_bf == 1)
                    {
                        RotateL(parent);
                    }
                    else if (cur->_bf == -1)
                    {
                        RotateRL(parent);
                    }
                }
                else if (parent->_bf == -2)
                {
                    if (cur->_bf == -1)
                    {
                        RotateR(parent);
                    }
                    else if (cur->_bf == 1)
                    {
                        RotateLR(parent);
                    }
                }
                else
                {
                    assert(false);
                    exit(-1);
                }
            }
            else
            {
                assert(false);
                exit(-1);
            }

        }
        return true;

    }

    template<class K, class V>
    void AVLtree<K, V>::RotateL(Node* parent)
    {
        Node* son = parent->_right;
        Node* sonL = son->_left;
        
        parent->_right = sonL;
        if (sonL != nullptr)
        {
            sonL->_parent = parent;
        }
        
        Node* granparent = parent->_parent;
        son->_parent = granparent;
        if (granparent != nullptr)
        {
            if (granparent->_left == parent)
            {
                granparent->_left = son;
            }
            else if (granparent->_right == parent)
            {
                granparent->_right = son;
            }
        }
        else
        {
            _root == son;
        }

        son->_left = parent;
        parent->_parent = son;

        //更新平衡因子
        parent->_bf = 0;
        son->_bf = 0;
    }

    template<class K, class V>
    void AVLtree<K, V>::RotateR(Node* parent)
    {
        Node* son = parent->_left;
        Node* sonR = son->_right;

        parent->_left = sonR;
        if (sonR != nullptr)
        {
            sonR->_parent = parent;
        }
        
        Node* granparent = parent->_parent;
        son->_parent = granparent;
        if (granparent != nullptr)
        {
            if (granparent->_left == parent)
            {
                granparent->_left = son;
            }
            else if (granparent->_right == parent)
            {
                granparent->_right = son;
            }
        }
        else
        {
            _root = son;
        }

        son->_right = parent;
        parent->_parent = son;

        son->_bf = parent->_bf = 0;
    }

    template<class K, class V>
    void AVLtree<K, V>::RotateRL(Node* parent)
    {
        Node* son = parent->_right;
        Node* sonL = son->_left;
        int bf = sonL->_bf;
        
        RotateR(parent->_right);
        RotateL(parent);

        if (bf == -1)
        {
            sonL->_bf = 0;
            parent->_bf = 1;
            son->_bf = 0;
        }
        else if (bf == 1)
        {
            son->_bf = -1;
            parent->_bf = sonL->_bf = 0;
        }
        else if (bf == 0)
        {
            son->_bf = parent->_bf = sonL->_bf = 0;
        }
        else
        {
            assert(false);
        }
    }

    template<class K, class V>
    void AVLtree<K, V>::RotateLR(Node* parent)
    {
        Node* son = parent->_left;
        Node* sonR = son->_right;
        int bf = sonR->_bf;
        
        RotateL(parent->_left);
        RotateR(parent);

        if (bf == -1)
        {
            son->_bf = 1;
            parent->_bf = sonR->_bf = 0;
        }
        else if (bf == 1)
        {
            parent->_bf = -1;
            son->_bf = sonR->_bf = 0;
        }
        else if (bf == 0)
        {
            son->_bf = parent->_bf = sonR->_bf = 0;
        }
        else
        {
            assert(false);
        }
    }

    template<class K, class V>
    void AVLtree<K, V>::_isBinarySerchTree(Node* root)
    {
        if (root == nullptr)
        {
            return;
        }

        _isBinarySerchTree(root->_left);

        std::cout << root->_kv.first << ' ';

        _isBinarySerchTree(root->_right);
    }

    template<class K, class V>
    bool AVLtree<K, V>::_isBalanceTree(Node* root)
    {
        if (nullptr == root)
        {
            return true;
        }

        if (root->_bf == 2 || root->_bf == -2)
        {
            return false;
        }

        bool left = _isBalanceTree(root->_left);
        bool right = _isBalanceTree(root->_right);

        bool cur = root->_bf == 2 || root->_bf == -2;

        return left && right && !cur;
    }

    template<class K, class V>
    int AVLtree<K, V>::Hight(Node* root)
    {
        if (nullptr == root)
        {
            return 0;
        }

        int left = Hight(root->_left);
        int right = Hight(root->_right);

        return left > right ? left + 1 : right + 1;
    }

    template<class K, class V>
    bool AVLtree<K, V>::_BalanceFactor(Node* root)
    {
        if (nullptr == root)
        {
            return true;
        }

        int leftHight = Hight(root->_left);
        int rightHight = Hight(root->_right);

        if (rightHight - leftHight != root->_bf)
        {
            std::cout << root->_kv.first << "平衡因子异常" <<  std::endl;
            return false;
        }

        return abs(rightHight - leftHight) < 2
            && _BalanceFactor(root->_left)
            && _BalanceFactor(root->_right);

    }

}